import os
import sqlite3
from datetime import datetime
from api.websocketapi import upload_image, get_images ,clean_dir
import websocket
from workflowprompt import color_flux as color_flux
import json
from config.host import server_address, client_id
from utils.comic import is_black_white_comic
from utils.image_utils import split_image, merge_images
from PIL import Image
import io
import copy
import shutil
import random
from tqdm import tqdm # 新增导入

# 可配置数据库路径
DB_DIR = 'E:\\software\\jmcomic_v1.2.6_windows_x64\\jmcomic\\cache'
#os.makedirs(DB_DIR, exist_ok=True)
db_path = os.path.join(DB_DIR, 'process_log.db')

conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute('''
CREATE TABLE IF NOT EXISTS process_log (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    src_name TEXT,
    out_name TEXT,
    dir_lv2 TEXT,
    dir_lv1 TEXT,
    status TEXT,
    msg TEXT,
    process_time TEXT
)
''')
# --- 新增：创建已处理文件夹记录表 ---
cursor.execute('''
CREATE TABLE IF NOT EXISTS processed_folders (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    folder_path TEXT UNIQUE,
    status TEXT,
    processed_time TEXT
)
''')
conn.commit()

# --- 修改：模型处理尺寸、重叠和缩放阈值配置 ---
TARGET_HEIGHT = 1568
OVERLAP = 64
RESIZE_THRESHOLD_RATIO = 1.3 # 当图片高度/目标高度小于此值时，缩放而非分割
SAMPLE_SIZE = 10 # 新增：随机抽样检查的图片数量
CONSECUTIVE_COLOR_TO_SKIP = 3 # 新增：连续多少张彩色图就跳过文件夹
MIN_BW_TO_PROCESS = 3 # 新增：至少有多少张黑白图才处理文件夹

if __name__ == '__main__':
    bookpath='E:\\software\\jmcomic_v1.2.6_windows_x64\\jmcomic\\cache\\book'
    img_exts = ['.jpg', '.jpeg', '.png', '.webp', '.bmp']

    # --- 新增：加载已处理的文件夹列表 ---
    processed_folders_set = set()
    try:
        cursor.execute("SELECT folder_path FROM processed_folders")
        for row in cursor.fetchall():
            processed_folders_set.add(row[0])
        print(f"--- 已加载 {len(processed_folders_set)} 个已处理/已跳过的文件夹记录 ---")
    except sqlite3.Error as e:
        print(f"加载已处理文件夹列表失败: {e}")


    # --- 新增：在处理前优先执行清理逻辑 ---
    print("--- 开始执行预处理清理 ---")
    for bookdir in os.listdir(bookpath):
        bookdir_path = os.path.join(bookpath, bookdir)
        if not os.path.isdir(bookdir_path):
            continue

        # 1. 检查并删除图片过少的章节文件夹
        if not os.listdir(bookdir_path): # 如果图书文件夹本身就是空的
             try:
                print(f"图书文件夹 '{bookdir}' 为空，正在删除...")
                os.rmdir(bookdir_path)
                print(f"已删除空图书文件夹: {bookdir_path}")
             except OSError as e:
                print(f"删除空图书文件夹失败: {bookdir_path}, 错误: {e}")
             continue

        for chapterdir in os.listdir(bookdir_path):
            chapterdir_path = os.path.join(bookdir_path, chapterdir)
            if not os.path.isdir(chapterdir_path):
                continue
            
            image_count = len([f for f in os.listdir(chapterdir_path) if any(f.lower().endswith(ext) for ext in img_exts)])
            
            if image_count < 3:
                try:
                    print(f"章节 '{chapterdir}' 图片数 ({image_count}) 小于3，正在删除...")
                    shutil.rmtree(chapterdir_path)
                    print(f"已删除章节文件夹: {chapterdir_path}")
                except OSError as e:
                    print(f"删除章节文件夹失败: {chapterdir_path}, 错误: {e}")

        # 2. 检查并删除已无任何章节的图书文件夹
        remaining_items = os.listdir(bookdir_path)
        if not any(os.path.isdir(os.path.join(bookdir_path, item)) for item in remaining_items):
            try:
                print(f"图书文件夹 '{bookdir}' 已无任何章节目录，正在删除...")
                shutil.rmtree(bookdir_path)
                print(f"已删除图书文件夹: {bookdir_path}")
            except OSError as e:
                print(f"删除图书文件夹失败: {bookdir_path}, 错误: {e}")
    print("--- 预处理清理完成 ---")


    # --- 开始正式的图片处理流程 ---
    # 获取所有图书目录及其修改时间
    book_dirs = [
        (bookdir, os.path.getmtime(os.path.join(bookpath, bookdir)))
        for bookdir in os.listdir(bookpath)
        if os.path.isdir(os.path.join(bookpath, bookdir))
    ]
    # 按修改时间倒序排序
    book_dirs.sort(key=lambda x: x[1], reverse=True)

    for bookdir, _ in book_dirs:
        bookdir_path = os.path.join(bookpath, bookdir)
        if not os.path.isdir(bookdir_path):
            continue
        for chapterdir in os.listdir(bookdir_path):
            chapterdir_path = os.path.join(bookdir_path, chapterdir)
            if not os.path.isdir(chapterdir_path):
                continue
            
            # --- 新增：检查文件夹是否需要跳过 ---
            if chapterdir_path in processed_folders_set:
                print(f"--- 章节 '{chapterdir}' 已被记录，本次跳过 ---")
                continue

            # --- 文件夹抽样检查逻辑 ---
            all_images_in_chapter = [f for f in os.listdir(chapterdir_path) if any(f.lower().endswith(ext) for ext in img_exts) and not f.endswith('_color' + os.path.splitext(f)[1])]
            
            if not all_images_in_chapter:
                continue

            random.shuffle(all_images_in_chapter)
            sample_images = all_images_in_chapter[:SAMPLE_SIZE]
            
            bw_count = 0
            consecutive_color_count = 0
            should_process_folder = False

            print(f"--- 对章节 '{chapterdir}' 进行抽样检查 (样本数: {len(sample_images)}) ---")
            for img_sample_name in sample_images:
                is_bw = is_black_white_comic(os.path.join(chapterdir_path, img_sample_name), color_ratio_threshold=0.5)
                if is_bw:
                    bw_count += 1
                    consecutive_color_count = 0 # 重置连续彩色计数
                else:
                    consecutive_color_count += 1
                
                # 检查是否满足连续彩色跳过条件
                if consecutive_color_count >= CONSECUTIVE_COLOR_TO_SKIP:
                    print(f"检测到连续 {CONSECUTIVE_COLOR_TO_SKIP} 张彩色图片，跳过此章节。")
                    should_process_folder = False
                    break
            
            # 在抽样结束后，根据黑白图片数量决定是否处理
            if not should_process_folder and bw_count > MIN_BW_TO_PROCESS:
                should_process_folder = True

            if not should_process_folder:
                print(f"抽样检查未通过 (黑白图片数: {bw_count})，跳过章节 '{chapterdir}'。")
                # --- 新增：记录跳过的文件夹 ---
                try:
                    cursor.execute(
                        "INSERT OR IGNORE INTO processed_folders (folder_path, status, processed_time) VALUES (?, ?, ?)",
                        (chapterdir_path, 'skipped_sampling', datetime.now().isoformat())
                    )
                    conn.commit()
                except sqlite3.Error as e:
                    print(f"记录跳过文件夹失败: {e}")
                continue
            
            print(f"抽样检查通过 (黑白图片数: {bw_count})，开始处理章节 '{chapterdir}'。")
            # --- 抽样检查结束 ---

            # --- 图片处理主循环 ---
            all_processed_successfully = True
            # 使用 tqdm 包装迭代器以显示进度条
            all_images_in_chapter.sort() 
            for imgname in tqdm(all_images_in_chapter, desc=f"处理章节 '{chapterdir}'", unit="张"):
                img_path = os.path.join(chapterdir_path, imgname)
                name, ext = os.path.splitext(imgname)
                
                # tqdm.write 可以避免破坏进度条的显示，但为了简单起见，保留 print 也可以
                print(f"处理图片: {img_path}") 
                
                # 已在抽样逻辑中过滤，但保留以防万一
                if name.endswith('_color') :
                    print(f"跳过已处理图片: {img_path}")
                    continue

                # 对单张图片进行最终的黑白确认，因为抽样可能不完全
                result = is_black_white_comic(img_path,color_ratio_threshold=0.5)
                if not result:
                    print(f"跳过非黑白漫画图片: {img_path}")
                    continue
                
                out_name = f"{name}_color{ext}"
                out_path = os.path.join(chapterdir_path, out_name)

                if os.path.exists(out_path):
                    print(f"已存在，跳过: {out_path}")
                    continue
                # 目录分级
                dir_lv1 = os.path.basename(os.path.dirname(img_path))
                dir_lv2 = os.path.basename(os.path.dirname(os.path.dirname(img_path)))

                # --- 修改后的处理逻辑 ---
                try:
                    # --- 新增：加载图片后，先进行黑白（灰度）处理 ---
                    # 1. 打开图片
                    img = Image.open(img_path)
                    # 2. 转换为灰度图以清除色彩噪点
                    grayscale_img = img.convert('L')
                    # 3. 转换回RGB模式以满足模型输入要求
                    original_img = grayscale_img.convert('RGB')
                    # original_img = Image.open(img_path)
                    original_width, original_height = original_img.size

                    final_image = None
                    
                    # 判断处理模式
                    mode = "standard"
                    if original_height > TARGET_HEIGHT:
                        ratio = original_height / TARGET_HEIGHT
                        if ratio <= RESIZE_THRESHOLD_RATIO:
                            mode = "resize"
                        else:
                            mode = "split"

                    # 模式1: 分割处理 (非常长的图)
                    if mode == "split":
                        print(f"长图分割模式: 分割图片 {imgname}")
                        image_blocks = split_image(img_path, block_height=TARGET_HEIGHT, overlap=OVERLAP)
                        processed_blocks = []
                        
                        for i, block in enumerate(image_blocks):
                            print(f"  处理块 {i+1}/{len(image_blocks)}...")
                            # 将块数据直接上传
                            buffer = io.BytesIO()
                            block.save(buffer, format='PNG')
                            buffer.seek(0)
                            upload_result = upload_image(buffer, filename=f"temp_chunk_{i}.png")

                            if upload_result['status'] != 200:
                                raise Exception(f"块上传失败: {upload_result['content']}")
                            
                            content = json.loads(upload_result['content'])
                            filename = os.path.basename(content['name'])

                            ws = websocket.WebSocket()
                            ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
                            prompt = copy.deepcopy(color_flux.prompt)
                            prompt['199']['inputs']['image'] = filename
                            images = get_images(ws, prompt)
                            ws.close()

                            processed_block_data = list(images.values())[0][0]
                            processed_block = Image.open(io.BytesIO(processed_block_data))
                            processed_blocks.append(processed_block)
                        
                        print("  拼接所有已处理的块...")
                        final_image = merge_images(processed_blocks, original_height, overlap=OVERLAP)

                    # 模式2: 缩放处理 (略长的图)
                    elif mode == "resize":
                        print(f"缩放模式: 处理图片 {imgname}")
                        # 缩放图片以适应目标高度
                        new_width = int(original_width * (TARGET_HEIGHT / original_height))
                        resized_img = original_img.resize((new_width, TARGET_HEIGHT), Image.LANCZOS)
                        
                        buffer = io.BytesIO()
                        resized_img.save(buffer, format='PNG')
                        buffer.seek(0)
                        upload_result = upload_image(buffer, filename=f"resized_{imgname}")

                        if upload_result['status'] != 200:
                            raise Exception(f"缩放图上传失败: {upload_result['content']}")
                        
                        content = json.loads(upload_result['content'])
                        filename = os.path.basename(content['name'])

                        ws = websocket.WebSocket()
                        ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
                        prompt = color_flux.prompt
                        prompt['199']['inputs']['image'] = filename
                        images = get_images(ws, prompt)
                        ws.close()

                        image_data = list(images.values())[0][0]
                        processed_image = Image.open(io.BytesIO(image_data))
                        # 将处理后的图片缩放回原始尺寸
                        #final_image = processed_image.resize((original_width, original_height), Image.LANCZOS)
                        final_image = processed_image

                    # 模式3: 标准处理 (高度合适的图)
                    else: # mode == "standard"
                        print(f"标准模式: 处理图片 {imgname}")
                        upload_result = upload_image(img_path)
                        if upload_result['status'] != 200:
                            raise Exception(f"上传失败: {upload_result['content']}")

                        content = json.loads(upload_result['content'])
                        filename = os.path.basename(content['name'])

                        ws = websocket.WebSocket()
                        ws.connect(f"ws://{server_address}/ws?clientId={client_id}")
                        prompt = color_flux.prompt
                        prompt['199']['inputs']['image'] = filename
                        images = get_images(ws, prompt)
                        ws.close()

                        image_data = list(images.values())[0][0]
                        final_image = Image.open(io.BytesIO(image_data))

                    # 保存最终图片并记录日志
                    if final_image:
                        final_image.save(out_path)
                        print(f"已保存: {out_path}")
                        cursor.execute(
                            "INSERT INTO process_log (src_name, out_name, dir_lv2, dir_lv1, status, msg, process_time) VALUES (?, ?, ?, ?, ?, ?, ?)",
                            (imgname, out_name, dir_lv2, dir_lv1, 'success', '', datetime.now().isoformat())
                        )
                        conn.commit()
                    
                    clean_result = clean_dir()
                    print(f"清理临时文件: {clean_result}")

                except Exception as e:
                    print(f"处理图片时出错: {imgname}, 错误: {e}")
                    cursor.execute(
                        "INSERT INTO process_log (src_name, out_name, dir_lv2, dir_lv1, status, msg, process_time) VALUES (?, ?, ?, ?, ?, ?, ?)",
                        (imgname, '', dir_lv2, dir_lv1, 'fail', str(e), datetime.now().isoformat())
                    )
                    conn.commit()
                    continue
            
            # --- 新增：在章节所有图片处理完后，记录文件夹状态 ---
            if all_processed_successfully:
                print(f"--- 章节 '{chapterdir}' 全部处理完成，记录状态 ---")
                try:
                    cursor.execute(
                        "INSERT OR IGNORE INTO processed_folders (folder_path, status, processed_time) VALUES (?, ?, ?)",
                        (chapterdir_path, 'completed', datetime.now().isoformat())
                    )
                    conn.commit()
                except sqlite3.Error as e:
                    print(f"记录完成文件夹失败: {e}")

conn.close()
